use std::{
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
    sync::Arc,
    time::Duration,
};

use tokio::{
    io,
    net::{lookup_host, TcpStream},
    sync::mpsc,
};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt, BufWriter},
    time::timeout,
};
use tracing::{info, trace, warn};

use crate::{consts, AddrType, AuthMethods, ReplyCode, Socks5Command, Socks5Req, SocksReply, User};

#[derive(Debug)]
pub struct Handler {
    // The `TcpStream`. It is decorated with a `BufWriter`, which provides write
    // level buffering. The `BufWriter` implementation provided by Tokio is
    // sufficient for our needs.
    pub stream: BufWriter<TcpStream>,
    client_addr: SocketAddr,
    users: Arc<Vec<User>>,
    auth_methods: Arc<Vec<u8>>,
    timeout: Option<Duration>,
    version: u8,
    auth_nmethods: u8,

    /// Listen for shutdown notifications.
    ///
    /// A wrapper around the `broadcast::Receiver` paired with the sender in
    /// `Listener`. The connection handler processes requests from the
    /// connection until the peer disconnects **or** a shutdown notification is
    /// received from `shutdown`. In the latter case, any in-flight work being
    /// processed for the peer is continued until it reaches a safe state, at
    /// which point the connection is terminated.
    // shutdown: Shutdown,

    /// Not used directly. Instead, when `Handler` is dropped...?
    _shutdown_complete: mpsc::Sender<()>,
}

impl Handler {
    /// Process a single connection.
    ///
    /// Request frames are read from the socket and processed. Responses are
    /// written back to the socket.
    ///
    /// Currently, pipelining is not implemented. Pipelining is the ability to
    /// process more than one request concurrently per connection without
    /// interleaving frames. See for more details:
    /// https://redis.io/topics/pipelining
    ///
    /// When the shutdown signal is received, the connection is processed until
    /// it reaches a safe state, at which point it is terminated.
    ///
    // #[instrument(skip(self))]
    // async fn run(&mut self) -> crate::Result<()> {
    //     // As long as the shutdown signal has not been received, try to read a
    //     // new request frame.
    //     while !self.shutdown.is_shutdown() {
    //         let mut header = [0u8, 2];
    //         self.stream.read_exact(&mut header).await?;

    //         let ver = header[0];
    //         let auth_methods = header[1];

    //         trace!("version: {}, auth methods: {}", ver, auth_methods);

    //         match ver {
    //             consts::SOCKS5_VERSION => {
    //                 // Authenticate client
    //                 self.auth().await?;
    //                 // Handle requests
    //                 self.handle_client().await?;
    //             }
    //             _ => {
    //                 warn!("unsupported socks version: {}", ver);
    //                 self.shutdown().await?;
    //             }
    //         }
    //     }
    //     Ok(())
    // }

    pub fn new(
        stream: TcpStream,
        client_addr: SocketAddr,
        authed_users: Arc<Vec<User>>,
        auth_methods: Arc<Vec<u8>>,
        timeout: Option<Duration>,
        // shutdown: Shutdown,
        shutdown_complete: mpsc::Sender<()>,
    ) -> Self {
        Handler {
            stream: BufWriter::new(stream),
            client_addr,
            users: authed_users,
            auth_methods: auth_methods,
            timeout: timeout,
            version: 0,
            auth_nmethods: 0,
            // shutdown: shutdown,
            _shutdown_complete: shutdown_complete,
        }
    }

    pub async fn init(&mut self) -> crate::Result<()> {
        let mut header = [0u8; 2];
        self.stream.read_exact(&mut header).await?;

        self.version = header[0];
        self.auth_nmethods = header[1];

        trace!(
            "[{:?}] version: {}, auth_nmethods: {}",
            self.client_addr,
            self.version,
            self.auth_nmethods
        );

        match self.version {
            consts::SOCKS5_VERSION => {
                // Authenticate w/ client
                self.auth().await?;
                // Handle requests
                self.handle_client().await?;
            }
            _ => {
                warn!("unsupported socks version: {} ", self.version);
                self.shutdown().await?;
            }
        }

        Ok(())
    }

    pub async fn auth(&mut self) -> crate::Result<()> {
        let methods = self.get_avalible_methods().await?;
        trace!("[{:?}] methods: {:?}", self.client_addr, methods);

        let mut response = [0u8, 2];
        response[0] = consts::SOCKS5_VERSION;

        if methods.contains(&(AuthMethods::UserPass as u8)) {
            response[1] = AuthMethods::UserPass as u8;

            info!("sending USER/PASS packet");
            // self.stream.write_all(&response).await?;
            // self.stream.flush().await?;

            self.write(&response.to_vec()).await?;

            let mut header = [0u8, 2];

            self.stream.read_exact(&mut header).await?;

            let ulen = header[1] as usize;

            let mut username = vec![0; ulen];

            self.stream.read_exact(&mut username).await?;

            let mut plen = [0u8; 1];
            self.stream.read_exact(&mut plen).await?;

            let mut password = vec![0; plen[0] as usize];
            self.stream.read_exact(&mut password).await?;

            let username = String::from_utf8_lossy(&username).to_string();
            let password = String::from_utf8_lossy(&password).to_string();

            let user = User { username, password };
            if self.authed(&user) {
                info!("access granted. user: {}", user.username);
                let response = [1, ReplyCode::Successed as u8];
                self.stream.write_all(&response).await?;
            } else {
                warn!("access denied. user: {}", user.username);
                let response = [1, ReplyCode::Failure as u8];
                self.stream.write_all(&response).await?;
                self.shutdown().await?;
            }
            Ok(())
        } else if methods.contains(&(AuthMethods::NoAuth as u8)) {
            response[1] = AuthMethods::NoAuth as u8;
            self.stream.write_all(&response).await?;
            Ok(())
        } else {
            response[1] = AuthMethods::NoMethods as u8;
            self.stream.write_all(&response).await?;
            self.shutdown().await?;
            Err("client has no suitable auth methods!".into())
        }
    }

    pub async fn handle_client(&mut self) -> crate::Result<usize> {
        let req = Socks5Req::from_stream(&mut self.stream).await?;

        match req.command {
            Socks5Command::TCPConnect => {
                let sock_addr = addr_to_socket(&req.addr_type, &req.addr, req.port).await?;

                let time_out = if let Some(timeout) = self.timeout {
                    timeout
                } else {
                    Duration::from_millis(500)
                };

                let mut target =
                    timeout(
                        time_out,
                        async move { TcpStream::connect(&sock_addr[..]).await },
                    )
                    .await
                    .map_err(|_| format!("connection refuesed"))??;

                SocksReply::new(ReplyCode::Successed)
                    .send(&mut self.stream)
                    .await?;

                match tokio::io::copy_bidirectional(&mut self.stream, &mut target).await {
                    Err(e) if e.kind() == std::io::ErrorKind::NotConnected => {
                        warn!("already closed");
                        Ok(0)
                    }
                    Err(e) => Err(e.into()),
                    Ok((_s_to_t, t_to_s)) => Ok(t_to_s as usize),
                }
            }
            Socks5Command::TCPBind => Err("bind not supported".into()),
            Socks5Command::UDPAssociate => Err("udp associate not supported".into()),
            _ => Err("unknown command".into()),
        }
    }

    /// Shutdown a client
    pub async fn shutdown(&mut self) -> io::Result<()> {
        info!("close connection, clinet_addr: {:?}", self.client_addr);
        self.stream.shutdown().await?;
        Ok(())
    }

    pub async fn write(&mut self, buf: &Vec<u8>) -> io::Result<()> {
        self.stream.write_all(buf).await?;
        self.stream.flush().await
    }

    /// Check if username + password pair are valid
    fn authed(&self, user: &User) -> bool {
        self.users.contains(user)
    }

    async fn get_avalible_methods(&mut self) -> crate::Result<Vec<u8>> {
        let mut methods: Vec<u8> = Vec::with_capacity(self.auth_nmethods as usize);
        for _ in 0..self.auth_nmethods {
            let mut method = [0u8; 1];
            self.stream.read_exact(&mut method).await?;
            if self.auth_methods.contains(&method[0]) {
                methods.append(&mut method.to_vec());
            }
        }
        Ok(methods)
    }
}

async fn addr_to_socket(
    addr_type: &AddrType,
    addr: &[u8],
    port: u16,
) -> io::Result<Vec<SocketAddr>> {
    match addr_type {
        AddrType::V6 => {
            let new_addr = (0..8)
                .map(|x| {
                    trace!("{} and {}", x * 2, (x * 2) + 1);
                    (u16::from(addr[x * 2]) << 8) | u16::from(addr[(x * 2) + 1])
                })
                .collect::<Vec<u16>>();

            Ok(vec![SocketAddr::from(SocketAddrV6::new(
                Ipv6Addr::new(
                    new_addr[0],
                    new_addr[1],
                    new_addr[2],
                    new_addr[3],
                    new_addr[4],
                    new_addr[5],
                    new_addr[6],
                    new_addr[7],
                ),
                port,
                0,
                0,
            ))])
        }
        AddrType::V4 => Ok(vec![SocketAddr::from(SocketAddrV4::new(
            Ipv4Addr::new(addr[0], addr[1], addr[2], addr[3]),
            port,
        ))]),
        AddrType::Domain => {
            let mut domain = String::from_utf8_lossy(addr).to_string();
            domain.push(':');
            domain.push_str(&port.to_string());

            Ok(lookup_host(domain).await?.collect())
        }
        _ => Ok(vec![SocketAddr::new(
            IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
            8080,
        )]),
    }
}
